import math

import torch

from problems.problem import Problem
from utils.vectors import get_ref_vectors


class DTLZ(Problem):
    def __init__(self, var_dim: int, obj_dim: int, max_fun_eval, kwargs: dict):
        low_limit = torch.zeros(1, var_dim)
        high_limit = torch.ones(1, var_dim)
        super().__init__(var_dim, obj_dim, low_limit, high_limit, max_fun_eval, kwargs)

    def get_g(self, dec):
        return 100 * (self.var_dim - self.obj_dim + 1 + torch.sum(
            (dec[:, self.obj_dim - 1:] - 0.5) ** 2 - torch.cos(20 * torch.pi * (dec[:, self.obj_dim - 1:] - 0.5)),
            dim=1)).unsqueeze(dim=1)

    def calc_zdt4_5(self, dec, g):
        cumprod_cos = torch.cumprod(
            torch.cat((torch.ones(g.size(0), 1), torch.cos(dec[:, :self.obj_dim - 1] * torch.pi / 2)), dim=1), dim=1)
        sin_term = torch.cat(
            (torch.ones(g.size(0), 1), torch.sin(torch.fliplr(dec[:, :self.obj_dim - 1]) * torch.pi / 2)), dim=1)
        return (1 + g).unsqueeze(1) * torch.fliplr(cumprod_cos) * sin_term

    def calc_zdt5_6(self, dec, g):
        temp = g.unsqueeze(1).repeat(1, self.obj_dim - 2)
        dec_temp = dec.clone()
        dec_temp[:, 1:self.obj_dim - 1] = (1 + 2 * temp * dec[:, 1:self.obj_dim - 1]) / (2 + 2 * temp)
        return self.calc_zdt4_5(dec_temp, g)

    def get_optimal_solutions0(self, size):
        return get_ref_vectors(size, self.obj_dim) / 2

    def get_optimal_solutions1(self, size):
        R = get_ref_vectors(size, self.obj_dim)
        norm = torch.sqrt(torch.sum(R ** 2, dim=1)).view(-1, 1)
        R = R / norm
        return R

    def get_optimal_solutions2(self, size):
        # 生成初始均匀分布的点
        R = torch.cat((torch.arange(0, 1 + 1 / (size - 1), 1 / (size - 1), dtype=torch.double).unsqueeze(1),
                       torch.arange(1, -1 / (size - 1), -1 / (size - 1), dtype=torch.double).unsqueeze(1)), dim=1)
        R = R / torch.sqrt(torch.sum(R ** 2, dim=1, keepdim=True))
        R = torch.cat((R[:, torch.zeros(self.obj_dim - 2, dtype=int)], R), dim=1)
        temp = [i for i in range(self.obj_dim - 2, -1, -1)]
        temp.insert(0, self.obj_dim - 2)
        power = torch.tensor(temp)
        return R / math.sqrt(2) ** power
